#ifndef CKKS_SRC_BOOTSTRAPPING_HOMOMORPHIC_ENCODING_H
#define CKKS_SRC_BOOTSTRAPPING_HOMOMORPHIC_ENCODING_H

#include <cmath>
#include <complex>
#include <vector>
#include <map>
#include <algorithm>

#include <vector>
#include "../scheme.h"


class EncodingMatrix {
public:
    Scheme &scheme;

    SecretKey &secretKey;

    vector<vector<Plaintext>> mU0hatTPreFFT; // CTS预处理矩阵

    vector<vector<vector<complex<double>>>> cts_yuan; // CTS未编码矩阵

    vector<vector<Plaintext>> mU0PreFFT; // STC预处理矩阵

    vector<int> mParamsEnc; // CTS参数

    vector<int> mParamsDec; // CTS参数

    // 构造函数
    EncodingMatrix(SecretKey &secretKey, Scheme &scheme);
    // 拷贝构造函数
    EncodingMatrix(const EncodingMatrix &cipher) = delete;
    EncodingMatrix &operator = (const EncodingMatrix &o) = delete;
    // 移动构造函数
    EncodingMatrix(EncodingMatrix &&other) = delete;
    EncodingMatrix &operator = (EncodingMatrix &&other) = delete;

    complex<double> *Rotate(const vector<complex<double>> &a, int index);

    // computes all powers of a primitive root of unity exp(2 * M_PI/m)
    vector<complex<double>> ComputeRoots(int N);

    // 生成矩阵分解的参数
    std::vector<int32_t> GetCollapsedFFTParams(uint32_t slots, uint32_t levelBudget, uint32_t dim1);

    /* *
     * Computes parameters to ensure the encoding and decoding computations take exactly the
     * specified number of levels. More specifically, it returns a vector than contains
     * layers (the number of layers to collapse in one level), rows (how many such levels),
     * rem (the number of layers remaining to be collapsed in one level)
     *
     * @param logSlots the base 2 logarithm of the number of slots.
     * @param budget the allocated level budget for the computation.
     */
    vector<uint32_t> SelectLayers(uint32_t logSlots, uint32_t budget);

    /* *
     * Computes the coefficients for the FFT encoding for CoeffEncodingCollapse such that every
     * iteration occupies one level.
     *
     * @param pows vector of roots of unity powers.
     * @param rotGroup rotation group indices to appropriately choose the elements of pows to compute iFFT.
     * @param flag_i flag that is 0 when we compute the coefficients for conj(U_0^T) and is 1 for conj(i*U_0^T).
     */
    std::vector<std::vector<std::complex<double>>> CoeffEncodingOneLevel(std::vector<std::complex<double>> pows,
                                                                         std::vector<uint32_t> rotGroup, bool flag_i);

    /* *
     * Computes the coefficients for the FFT decoding for CoeffDecodingCollapse such that every
     * iteration occupies one level.
     *
     * @param pows vector of roots of unity powers.
     * @param rotGroup rotation group indices to appropriately choose the elements of pows to compute iFFT.
     * @param flag_i flag that is 0 when we compute the coefficients for U_0 and is 1 for i*U_0.
     */
    std::vector<std::vector<std::complex<double>>> CoeffDecodingOneLevel(std::vector<std::complex<double>> pows,
                                                                         std::vector<uint32_t> rotGroup, bool flag_i);

    /* *
     * Computes the coefficients for the given level budget for the FFT encoding. Needed in
     * EvalLTFFTPrecomputeEncoding.
     *
     * @param pows vector of roots of unity powers.
     * @param rotGroup rotation group indices to appropriately choose the elements of pows to compute iFFT.
     * @param levelBudget the user specified budget for levels.
     * @param flag_i flag that is 0 when we compute the coefficients for conj(U_0^T) and is 1 for conj(i*U_0^T).
     */
    std::vector<std::vector<std::vector<std::complex<double>>>> CoeffEncodingCollapse(
            std::vector<std::complex<double>> pows, std::vector<uint32_t> rotGroup, uint32_t levelBudget, bool flag_i);

    /* *
     * Computes the coefficients for the given level budget for the FFT decoding. Needed in
     * EvalLTFFTPrecomputeDecoding.
     *
     * @param pows vector of roots of unity powers.
     * @param rotGroup rotation group indices to appropriately choose the elements of pows to compute FFT.
     * @param levelBudget the user specified budget for levels.
     * @param flag_i flag that is 0 when we compute the coefficients for U_0 and is 1 for i*U_0.
     */
    std::vector<std::vector<std::vector<std::complex<double>>>> CoeffDecodingCollapse(
            std::vector<std::complex<double>> pows, std::vector<uint32_t> rotGroup, uint32_t levelBudget, bool flag_i);

    std::vector<std::vector<Plaintext>> EvalCoeffsToSlotsPrecompute(const std::vector<std::complex<double>> &A,
                                                                    const std::vector<uint32_t> &rotGroup, bool flag_i, double scale = 1, uint32_t L = 0);

    std::vector<std::vector<Plaintext>> EvalSlotsToCoeffsPrecompute(const std::vector<std::complex<double>> &A,
                                                                    const std::vector<uint32_t> &rotGroup, bool flag_i, double scale, uint32_t L);


    uint32_t ReduceRotation(int32_t index, uint32_t slots);

    void MulScalarP(uint64_t *res, uint64_t *a, long l);

    vector<uint64_t *> EvalFastRotationPrecompute(Ciphertext &cipher);

    void EvalMultExt(Ciphertext &result, Ciphertext &cipher, Plaintext pt);

    Ciphertext EvalMultExt(Ciphertext &cipher, Plaintext pt);

    void EvalAddExtInPlace(Ciphertext &cipher, Ciphertext &ext);

    uint64_t *InnerProduct(int levelQ, vector<uint64_t *> aPolyTs, vector<vector<vector<uint64_t>>> key); // return to
    // PQ

    void AddAndEqual(uint64_t *a, uint64_t *b, long l, long k);

    void NTTAndEqual(uint64_t *a, long l, long k);

    void PermuteNTTWithIndex(uint64_t *polOut, uint64_t *polIn, uint64_t *index, int l, int k);

    void EvalFastRotationExt(Ciphertext &result, Ciphertext &cipher, vector<uint64_t *> digits, int32_t index,
                             bool addFirst);

    // 对密文进行模降QP-->Q
    // input   cipher(mod QP)
    // output  cipher(mod Q)
    Ciphertext KeySwitchDown(Ciphertext &cipher);

    Ciphertext KeySwitchDownFirst(Ciphertext &cipher);

    // 对密文进行模升Q-->QP
    // input   cipher(mod Q)
    // output  cipher(mod QP)
    void KeySwitchExt(Ciphertext &result, Ciphertext cipher);

    void PQDecrypt(SecretKey &secretKeyPQ, Ciphertext &cipher, string s);

    void PQRescaleAndDecrypt(SecretKey &secretKeyPQ, Ciphertext &cipher, string s);

    void QDecrypt(SecretKey &secretKey, Ciphertext &cipher, string s);

    void QRescaleAndDecrypt(SecretKey &secretKeyQ, Ciphertext &cipher, string s);

    Ciphertext EvalCoeffsToSlots(const vector<vector<Plaintext>> &A, Ciphertext &cipher);

    Ciphertext EvalSlotsToCoeffs(const vector<vector<Plaintext>> &A, Ciphertext &cipher);
};

#endif // CKKS_HOMOMORPHIC_ENCODING_H
